Added kl_divergence for multivariate normals #1654
Conversation
lumip
force-pushed
the
mvn-kl-divergence
branch
from
4d999ca to
304cab7
Compare
lumip
force-pushed
the
mvn-kl-divergence
branch
from
304cab7 to
4e3beae
Compare
making the linter tests happy
|
Thanks for the feedback @fehiepsi . I have modified the code according to your suggestions and also added a test case where the batch dimensions of the two distributions are not identical (and fixed the linter complaints). If everything is good now, I'd still rebase onto the recent master and merge the commits. |
numpyro/distributions/kl.py
Outdated
|
|
||
| Lq_inv = solve_triangular(q_scale_tril, jnp.eye(D), lower=True) | ||
| q_half_log_det = jnp.log(jnp.diagonal(q.scale_tril, axis1=-2, axis2=-1)).sum(-1) | ||
| assert q_half_log_det.shape == q.batch_shape |
There was a problem hiding this comment.
nit: this assertion might not be true. In MultivariateNormal implementation, we avoid unnecessary broadcasting (e.g. we can have a batch of means with a single scale_tril).
There was a problem hiding this comment.
Right, I wasn't thinking of that. Removed the assertion and added tests for those cases.
numpyro/distributions/kl.py
Outdated
| f" {p.event_shape} and {q.event_shape} for p and q, respectively." | ||
| ) | ||
|
|
||
| if p.batch_shape != q.batch_shape: | ||
| min_batch_ndim = min(len(p.batch_shape), len(q.batch_shape)) | ||
| if p.batch_shape[-min_batch_ndim:] != q.batch_shape[-min_batch_ndim:]: |
There was a problem hiding this comment.
how about only assert that p.batch_shape and q.batch_shape can be broadcasted.
try:
result_batch_shape = jnp.broadcast_shapes(p.batch_shape, q.batch_shape)
except ValueError:
raise ...
numpyro/distributions/kl.py
Outdated
| assert jnp.ndim(q_mean) == 1 | ||
| assert jnp.ndim(p_scale_tril) == 2 | ||
| assert jnp.ndim(q_scale_tril) == 2 | ||
| assert q.mean.shape == q.batch_shape + q.event_shape |
There was a problem hiding this comment.
those assertions are unnecessary.
numpyro/distributions/kl.py
Outdated
|
|
||
| return .5 * (tr + t1 - D - log_det_ratio) | ||
| tr = _batch_trace_from_cholesky(Lq_inv @ p.scale_tril) | ||
| assert tr.shape == result_batch_shape |
There was a problem hiding this comment.
this assertion might not be true.
numpyro/distributions/kl.py
Outdated
| p_mean_flat = jnp.reshape(p.mean, (-1, D)) | ||
| p_scale_tril_flat = jnp.reshape(p.scale_tril, (-1, D, D)) | ||
| t1 = jnp.square(Lq_inv @ (p.loc - q.loc)[..., jnp.newaxis]).sum((-2, -1)) | ||
| assert t1.shape == result_batch_shape |
There was a problem hiding this comment.
this assertion might not be true
numpyro/distributions/kl.py
Outdated
| ).sum(-1) | ||
| log_det_ratio = 2 * (p_half_log_det - q_half_log_det) | ||
| p_half_log_det = jnp.log(jnp.diagonal(p.scale_tril, axis1=-2, axis2=-1)).sum(-1) | ||
| assert p_half_log_det.shape == p.batch_shape |
There was a problem hiding this comment.
this assertion might not be true
test/test_distributions.py
Outdated
| ((), ()), | ||
| ((1,), (1,)), | ||
| ((2, 3), (2, 3)), | ||
| ((5, 2, 3), (2, 3)), |
There was a problem hiding this comment.
could you change this to (5, 1, 3) and (2, 3)?
test/test_distributions.py
Outdated
| ((1,), (1,)), | ||
| ((2, 3), (2, 3)), | ||
| ((5, 2, 3), (2, 3)), | ||
| ((2, 3), (5, 2, 3)), |
There was a problem hiding this comment.
nit: maybe ((1, 3), (5, 2, 3))?
KL implementation for
MultivariateNormalEDIT: just saw that there is already a pull request for this from a year ago (#1487). What's the hold up with that?